#include "grid_ground_filter.hpp"

#include <pcl/filters/extract_indices.h>
#include <pcl_conversions/pcl_conversions.h>

#include <cmath>
#include <algorithm>

namespace autoware_ground_segmentation
{

GridGroundFilter::GridGroundFilter()
{
  // Initialize grid parameters
  grid_size_m_ = 0.1;
  grid_mode_switch_radius_ = 20.0;
  gnd_grid_buffer_size_ = 4;
  virtual_lidar_z_ = 2.0;
}

GridGroundFilter::~GridGroundFilter() = default;

void GridGroundFilter::setParameters(float grid_size, float switch_radius, uint16_t buffer_size, float lidar_z)
{
  grid_size_m_ = grid_size;
  grid_mode_switch_radius_ = switch_radius;
  gnd_grid_buffer_size_ = buffer_size;
  virtual_lidar_z_ = lidar_z;
}

void GridGroundFilter::processGrid(
  const pcl::PointCloud<pcl::PointXYZ>::Ptr& input_cloud,
  pcl::PointIndices& ground_indices,
  pcl::PointIndices& no_ground_indices)
{
  ground_indices.indices.clear();
  no_ground_indices.indices.clear();
  
  if (input_cloud->empty()) {
    return;
  }
  
  // Create grid structure
  std::map<std::pair<int, int>, std::vector<size_t>> grid_map;
  
  // Assign points to grid cells
  for (size_t i = 0; i < input_cloud->size(); ++i) {
    const auto& point = input_cloud->points[i];
    
    // Calculate grid coordinates
    int grid_x = static_cast<int>(point.x / grid_size_m_);
    int grid_y = static_cast<int>(point.y / grid_size_m_);
    
    // Add point to grid cell
    grid_map[{grid_x, grid_y}].push_back(i);
  }
  
  // Process each grid cell
  for (const auto& grid_cell : grid_map) {
    const auto& point_indices = grid_cell.second;
    
    if (point_indices.empty()) continue;
    
    // Find lowest point in grid cell
    size_t lowest_idx = point_indices[0];
    float lowest_z = input_cloud->points[lowest_idx].z;
    
    for (size_t idx : point_indices) {
      if (input_cloud->points[idx].z < lowest_z) {
        lowest_z = input_cloud->points[idx].z;
        lowest_idx = idx;
      }
    }
    
    // Classify points based on height difference from lowest point
    for (size_t idx : point_indices) {
      const auto& point = input_cloud->points[idx];
      float height_diff = point.z - lowest_z;
      
      // Simple height-based classification
      if (height_diff < 0.2) { // 20cm threshold
        ground_indices.indices.push_back(idx);
      } else {
        no_ground_indices.indices.push_back(idx);
      }
    }
  }
}

void GridGroundFilter::processElevationGrid(
  const pcl::PointCloud<pcl::PointXYZ>::Ptr& input_cloud,
  pcl::PointIndices& ground_indices,
  pcl::PointIndices& no_ground_indices)
{
  ground_indices.indices.clear();
  no_ground_indices.indices.clear();
  
  if (input_cloud->empty()) {
    return;
  }
  
  // Create elevation grid
  std::map<std::pair<int, int>, std::vector<std::pair<float, size_t>>> elevation_grid;
  
  // Assign points to elevation grid
  for (size_t i = 0; i < input_cloud->size(); ++i) {
    const auto& point = input_cloud->points[i];
    
    // Calculate grid coordinates
    int grid_x = static_cast<int>(point.x / grid_size_m_);
    int grid_y = static_cast<int>(point.y / grid_size_m_);
    
    // Add point with height to grid cell
    elevation_grid[{grid_x, grid_y}].emplace_back(point.z, i);
  }
  
  // Process each grid cell
  for (auto& grid_cell : elevation_grid) {
    auto& height_points = grid_cell.second;
    
    if (height_points.empty()) continue;
    
    // Sort points by height
    std::sort(height_points.begin(), height_points.end(),
              [](const std::pair<float, size_t>& a, const std::pair<float, size_t>& b) {
                return a.first < b.first;
              });
    
    // Use lowest points as ground reference
    size_t ground_count = std::min(static_cast<size_t>(gnd_grid_buffer_size_), height_points.size());
    
    for (size_t i = 0; i < height_points.size(); ++i) {
      size_t point_idx = height_points[i].second;
      
      if (i < ground_count) {
        // Lower points are ground
        ground_indices.indices.push_back(point_idx);
      } else {
        // Higher points are non-ground
        no_ground_indices.indices.push_back(point_idx);
      }
    }
  }
}

} // namespace autoware_ground_segmentation 